import numpy as np
import matplotlib.pyplot as plt


class PIDController:
    def __init__(self, kp, ki, kd, max_output, dt):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.max_output = max_output
        self.dt = dt
        self.integral = 0.0
        self.prev_error = 0.0

    def update(self, setpoint, measurement):
        error = setpoint - measurement
        self.integral += error * self.dt
        derivative = (error - self.prev_error) / self.dt
        output = self.kp * error + self.ki * self.integral + self.kd * derivative
        output = np.clip(output, -self.max_output, self.max_output)
        self.prev_error = error
        return output


class PositionController:
    """位置环：生成期望推力、姿态角和偏航角"""

    def __init__(self, dt):
        # XYZ位置PID
        self.x_pid = PIDController(1.2, 0.05, 0.8, 2.0, dt)  # 输出俯仰角
        self.y_pid = PIDController(1.2, 0.05, 0.8, 2.0, dt)  # 输出横滚角
        self.z_pid = PIDController(15.0, 2.0, 5.0, 20.0, dt)  # 输出总推力

    def update(self, pos_set, pos_meas, vel_meas):
        # 水平位置控制转换为姿态角指令（弧度）
        theta_d = self.x_pid.update(pos_set[0], pos_meas[0])
        phi_d = -self.y_pid.update(pos_set[1], pos_meas[1])

        # 高度控制（需要重力补偿）
        f_d = self.z_pid.update(pos_set[2], pos_meas[2]) + 9.81  # 重力补偿

        return phi_d, theta_d, f_d  # 返回姿态角和总推力


class AttitudeController:
    """姿态环：生成期望力矩"""

    def __init__(self, dt):
        self.phi_pid = PIDController(8.0, 0.1, 2.5, 5.0, dt)  # 横滚
        self.theta_pid = PIDController(8.0, 0.1, 2.5, 5.0, dt)  # 俯仰
        self.psi_pid = PIDController(4.0, 0.05, 1.0, 2.0, dt)  # 偏航

    def update(self, att_set, att_meas, rates_meas):
        # att_set: [phi_d, theta_d, psi_d]
        tau_phi = self.phi_pid.update(att_set[0], att_meas[0])
        tau_theta = self.theta_pid.update(att_set[1], att_meas[1])
        tau_psi = self.psi_pid.update(att_set[2], att_meas[2])
        return np.array([tau_phi, tau_theta, tau_psi])


class ControlAllocator:
    """控制分配：将推力/力矩转换为电机指令"""

    def __init__(self):
        # 混控矩阵参数（示例为X型布局）
        self.mix_matrix = np.array([
            [1, 1, 1, 1],  # 总推力
            [-1, -1, 1, 1],  # 横滚力矩
            [-1, 1, 1, -1],  # 俯仰力矩
            [1, -1, 1, -1]  # 偏航力矩
        ])

    def allocate(self, f_d, tau_d):
        # 解混控方程：motor_cmds = inv(mix_matrix) * [f_d, tau_x, tau_y, tau_z]
        # 简化解算（需根据实际布局调整）
        cmd_base = f_d / 4
        m1 = cmd_base + (tau_d[0] - tau_d[1] + tau_d[2]) / 4
        m2 = cmd_base + (-tau_d[0] + tau_d[1] - tau_d[2]) / 4
        m3 = cmd_base + (tau_d[0] + tau_d[1] + tau_d[2]) / 4
        m4 = cmd_base + (-tau_d[0] - tau_d[1] - tau_d[2]) / 4
        return np.clip([m1, m2, m3, m4], 0, 1)


class QuadcopterDynamics:
    """四旋翼动力学模型"""

    def __init__(self, dt):
        # 状态变量
        self.position = np.zeros(3)  # x,y,z (m)
        self.velocity = np.zeros(3)
        self.attitude = np.zeros(3)  # phi,theta,psi (rad)
        self.rates = np.zeros(3)  # p,q,r (rad/s)

        # 物理参数
        self.mass = 1.0  # kg
        self.I = np.array([0.1, 0.1, 0.2])  # 转动惯量
        self.g = 9.81
        self.dt = dt

    def update(self, motor_cmds):
        # 电机推力模型（简化：指令直接映射为推力）
        thrusts = motor_cmds * 20.0  # 假设每个电机最大推力20N

        # 计算总力和力矩
        f_total = np.sum(thrusts)
        tau = np.array([
            (thrusts[3] + thrusts[2] - thrusts[1] - thrusts[0]) * 0.1,  # 横滚
            (thrusts[3] + thrusts[0] - thrusts[2] - thrusts[1]) * 0.1,  # 俯仰
            (thrusts[0] + thrusts[2] - thrusts[1] - thrusts[3]) * 0.2  # 偏航
        ])

        # 平移运动
        acceleration = np.array([0.0, 0.0, -self.g])
        acceleration[2] += f_total / self.mass  # Z轴加速度
        self.velocity += acceleration * self.dt
        self.position += self.velocity * self.dt

        # 旋转运动
        angular_acc = tau / self.I
        self.rates += angular_acc * self.dt
        self.attitude += self.rates * self.dt


# 仿真参数
dt = 0.01
sim_time = 10.0
steps = int(sim_time / dt)

# 初始化系统
quad = QuadcopterDynamics(dt)
pos_controller = PositionController(dt)
att_controller = AttitudeController(dt)
allocator = ControlAllocator()

# 设定值 [x, y, z, psi]
setpoint = np.array([2.0, 1.5, 5.0, 0.0])

# 数据记录
history = {
    'time': np.zeros(steps),
    'position': np.zeros((steps, 3)),
    'attitude': np.zeros((steps, 3))
}

# 主控制循环
for i in range(steps):
    # 位置控制（外环）
    phi_d, theta_d, f_d = pos_controller.update(
        setpoint[:3], quad.position, quad.velocity
    )

    # 姿态控制（内环）
    tau_d = att_controller.update(
        [phi_d, theta_d, setpoint[3]],  # 姿态设定值
        quad.attitude,
        quad.rates
    )

    # 控制分配
    motor_cmds = allocator.allocate(f_d, tau_d)

    # 更新动力学
    quad.update(motor_cmds)

    # 记录数据
    history['time'][i] = i * dt
    history['position'][i] = quad.position.copy()
    history['attitude'][i] = np.rad2deg(quad.attitude)

# 可视化
plt.figure(figsize=(12, 8))

# 位置响应
plt.subplot(2, 1, 1)
plt.plot(history['time'], history['position'])
plt.legend(['X', 'Y', 'Z'])
plt.ylabel('Position (m)')

# 姿态响应
plt.subplot(2, 1, 2)
plt.plot(history['time'], history['attitude'])
plt.legend(['Roll', 'Pitch', 'Yaw'])
plt.ylabel('Attitude (deg)')
plt.xlabel('Time (s)')

plt.tight_layout()
plt.show()